import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm1d(hidden_size)
        self.dropout1 = nn.Dropout(p=0.1)
        self.lstm1 = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
        self.dropout2 = nn.Dropout(p=0.1)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.lstm2 = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
        self.dropout3 = nn.Dropout(p=0.1)
        self.fc3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = out.transpose(1, 2)
        out = self.bn(out)
        out = out.transpose(1, 2)
        out = self.dropout1(out)
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()
        out, (hn, cn) = self.lstm1(out, (h0.detach(), c0.detach()))
        out = self.dropout2(out)
        out = self.fc2(out[:, -1, :])
        out = out.unsqueeze(1).repeat(1, x.size(1), 1)
        out, (hn, cn) = self.lstm2(out, (h0.detach(), c0.detach()))
        out = self.dropout3(out)
        out = self.fc3(out[:, -1, :])
        return out

